package com.example.spider.task.limit;

import com.example.spider.task.AbstractSpiderTask;
import com.example.spider.task.SpiderTaskQueue;
import com.example.spider.task.Task;
import org.shoulder.core.concurrent.Threads;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Semaphore;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;

/**
 * 用于控制最大同时并发数目：如最大 10TPS每秒，即：支持瞬间达到最大限制，但达到后需要有请求响应才能发下一个请求
 * 对服务器压力随服务器性能变化而变化
 *
 * @author lym
 */
public class MaxCurrentRequestLimit extends AbstractLimitStrategy {

    private int requestNum;

    private static final Map<String, Semaphore> semaphoreMap = new ConcurrentHashMap<>(16);

    public MaxCurrentRequestLimit(Map<String, Object> map) {
        // only for dejson
        super(map);
        this.requestNum = Integer.parseInt((String) map.get("requestNum"));
    }

    public MaxCurrentRequestLimit(Duration limitDuration, int requestNum) {
        super(limitDuration);
        this.requestNum = requestNum;
    }

    public MaxCurrentRequestLimit(String key, Duration limitDuration, LimitScope limitScope, int requestNum) {
        super(key, limitDuration, limitScope);
        this.requestNum = requestNum;
    }

    // ===================================================================================

    @Override
    public boolean needLimit(Task<?> task) {
        checkProperties();
        key = genLimitKey(task);

        return !semaphoreMap.computeIfAbsent(key, k -> new Semaphore(requestNum, false)).tryAcquire();
    }

    @Override
    public void doLimit(AbstractSpiderTask<?> task) {
        // 延迟放回队列重试，这里不要用阻塞方法，爬虫网络返回往往比较慢
        long waitTime = ThreadLocalRandom.current().nextLong(1, 100);
        Threads.delay(() -> SpiderTaskQueue.putTask(task), waitTime, TimeUnit.MILLISECONDS);
    }

    @Override
    public <T> void afterRequest(Task<T> task, T result, Throwable ex, long endTime) {
        if (ex == null) {
            semaphoreMap.computeIfAbsent(key, k -> new Semaphore(requestNum, false)).release();
        }
    }

    @Override
    public LimitStrategy copyWithNewKey(String s) {
        return new MaxCurrentRequestLimit(getKey(), getLimitDuration(), getLimitScope(), getRequestNum());
    }

    public int getRequestNum() {
        return requestNum;
    }

    public void setRequestNum(int requestNum) {
        this.requestNum = requestNum;
    }
}
